use std::borrow::Cow;
use std::cell::RefCell;
use std::collections::{BTreeMap, HashSet};
use std::fmt::Display;
use std::rc::Rc;

use log::debug;

use crate::c::{
    BareNumberC, CompimeConstant, CustomTypeDeclaration, CustomTypeDefinition, ExprC, FuncDeclare,
    StatPosC,
};
use crate::char_seq::Pos;
use crate::lexer::{Arithmetic, CompareOperator};
use crate::parser::Error;

type CowStr = Cow<'static, str>;

#[derive(Clone, Debug)]
pub struct Env(Rc<RefCell<EnvMap>>);

#[derive(Debug)]
struct EnvMap {
    constvar_map: BTreeMap<CowStr, Expr>,
    variable_type: BTreeMap<CowStr, Type>,
    struct_map: BTreeMap<CowStr, Struct>,
    enum_map: BTreeMap<CowStr, Enum>,
    union_map: BTreeMap<CowStr, Union>,
    parent: Option<Env>,
}

impl Env {
    pub fn new() -> Self {
        Env(Rc::new(RefCell::new(EnvMap {
            constvar_map: BTreeMap::new(),
            variable_type: BTreeMap::new(),
            struct_map: BTreeMap::new(),
            enum_map: BTreeMap::new(),
            union_map: BTreeMap::new(),
            parent: None,
        })))
    }

    pub fn with_parent(parent: Env) -> Self {
        Env(Rc::new(RefCell::new(EnvMap {
            constvar_map: BTreeMap::new(),
            variable_type: BTreeMap::new(),
            struct_map: BTreeMap::new(),
            enum_map: BTreeMap::new(),
            union_map: BTreeMap::new(),
            parent: Some(parent),
        })))
    }

    fn find_struct(&self, name: &str, pos: &Pos) -> Result<Struct, Error> {
        let env = self.0.as_ref().borrow();

        let ty = match env.struct_map.get(name) {
            Some(some) => some.clone(),
            None => match env.parent.clone() {
                Some(parent) => parent.find_struct(name, pos)?,
                None => {
                    return Err(Error::Syntax {
                        error: format!("unbound type '{name}'"),
                        pos: pos.clone(),
                    });
                }
            },
        };

        Ok(ty)
    }

    fn find_enum(&self, name: &str, pos: &Pos) -> Result<Enum, Error> {
        let env = self.0.as_ref().borrow();

        let ty = match env.enum_map.get(name) {
            Some(some) => some.clone(),
            None => match env.parent.clone() {
                Some(parent) => parent.find_enum(name, pos)?,
                None => {
                    return Err(Error::Syntax {
                        error: format!("unbound enum type '{name}'"),
                        pos: pos.clone(),
                    });
                }
            },
        };

        Ok(ty)
    }

    fn find_union(&self, name: &str, pos: &Pos) -> Result<Union, Error> {
        let env = self.0.as_ref().borrow();

        let ty = match env.union_map.get(name) {
            Some(some) => some.clone(),
            None => match env.parent.clone() {
                Some(parent) => parent.find_union(name, pos)?,
                None => {
                    return Err(Error::Syntax {
                        error: format!("unbound union type '{name}'"),
                        pos: pos.clone(),
                    });
                }
            },
        };

        Ok(ty)
    }

    fn find_const(&self, name: &str, pos: &Pos) -> Result<Expr, Error> {
        let env = self.0.as_ref().borrow();

        let ty = match env.constvar_map.get(name) {
            Some(some) => some.clone(),
            None => match env.parent.clone() {
                Some(parent) => parent.find_const(name, pos)?,
                None => {
                    return Err(Error::Syntax {
                        error: format!("unbound variable {name}"),
                        pos: pos.clone(),
                    });
                }
            },
        };

        Ok(ty)
    }

    fn find_var_type(&self, name: &str, pos: &Pos) -> Result<Type, Error> {
        let env = self.0.as_ref().borrow();

        let ty = match env.variable_type.get(name) {
            Some(some) => some.clone(),
            None => match env.parent.clone() {
                Some(parent) => parent.find_var_type(name, pos)?,
                None => {
                    return Err(Error::Syntax {
                        error: format!("unbound variable {name}"),
                        pos: pos.clone(),
                    });
                }
            },
        };

        Ok(ty)
    }

    fn bind_var_type(&self, name: CowStr, ty: Type, pos: &Pos) -> Result<(), Error> {
        let mut envmap = self.0.as_ref().borrow_mut();

        if envmap.variable_type.insert(name.clone(), ty).is_some() {
            return Err(Error::DuplicateName {
                name,
                pos: pos.clone(),
            });
        }

        Ok(())
    }

    fn bind_const(&self, name: CowStr, expr: Expr, pos: &Pos) -> Result<(), Error> {
        let mut envmap = self.0.as_ref().borrow_mut();

        if envmap.constvar_map.insert(name.clone(), expr).is_some() {
            return Err(Error::DuplicateName {
                name,
                pos: pos.clone(),
            });
        }
        Ok(())
    }

    fn bind_struct(&self, name: CowStr, ty: Struct) {
        let mut envmap = self.0.as_ref().borrow_mut();

        debug!("bind struct, name: {name}, ty: {ty:?}");

        envmap.struct_map.insert(name.clone(), ty);
    }

    fn bind_enum(&self, name: CowStr, ty: Enum) {
        let mut envmap = self.0.as_ref().borrow_mut();

        debug!("bind enum, name: {name}, ty: {ty:?}");

        envmap.enum_map.insert(name, ty);
    }

    fn bind_union(&self, name: CowStr, ty: Union) {
        let mut envmap = self.0.as_ref().borrow_mut();

        debug!("bind union, name: {name}, ty: {ty:?}");

        envmap.union_map.insert(name, ty);
    }
}

pub fn load_stdio(env: &Env) -> Result<(), Error> {
    env.bind_var_type(
        "puts".into(),
        Type::Func {
            return_type: Box::new(Type::Void),
            paras: vec![Type::Char.to_const().to_pointer()],
        },
        &Pos::zero(),
    )?;
    env.bind_var_type(
        "getchar".into(),
        Type::Func {
            return_type: Box::new(Type::Char),
            paras: vec![],
        },
        &Pos::zero(),
    )?;
    Ok(())
}

pub fn load_assert(env: &Env) -> Result<(), Error> {
    env.bind_var_type(
        "assert".into(),
        Type::Func {
            return_type: Box::new(Type::Void),
            paras: vec![Type::Boolean],
        },
        &Pos::zero(),
    )
}

pub fn load_stdlib(env: &Env) -> Result<(), Error> {
    env.bind_var_type(
        "malloc".into(),
        Type::Func {
            return_type: Box::new(Type::Void.to_pointer()),
            paras: vec![Type::UInt],
        },
        &Pos::zero(),
    )?;

    env.bind_var_type(
        "realloc".into(),
        Type::Func {
            return_type: Box::new(Type::Void.to_pointer()),
            paras: vec![Type::Void.to_optional_pointer(), Type::UInt],
        },
        &Pos::zero(),
    )?;

    env.bind_var_type(
        "free".into(),
        Type::Func {
            return_type: Box::new(Type::Void),
            paras: vec![Type::Void.to_optional_pointer()],
        },
        &Pos::zero(),
    )?;

    Ok(())
}

#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    Int(isize),
    UInt(usize),
    I32(i32),
    U32(u32),
    I64(i64),
    U64(u64),
    Null,

    Boolean(bool),
    String(CowStr),
    Char(char),

    Variable(CowStr),

    // NAME::VARIANT
    EnumGet {
        name: CowStr,
        variant: CowStr,
    },

    // EXPR.NAME
    // (. EXPR NAME)
    Get {
        expr: Box<ExprPos>,
        name: CowStr,
    },

    // EXPR[NAME]
    // ([] EXPR NAME)
    Take {
        expr: Box<ExprPos>,
        index: Box<ExprPos>,
    },

    // (cast TY EXPR)
    Cast {
        ty: Type,
        expr: Box<ExprPos>,
    },

    // (sizeof TY)
    Sizeof {
        ty: Type,
    },

    // (+ head rest...)
    Compute {
        op: Arithmetic,
        head: Box<ExprPos>,
        rest: Vec<ExprPos>,
    },

    // (&& head rest...)
    And {
        head: Box<ExprPos>,
        rest: Vec<ExprPos>,
    },

    // (|| head rest...)
    Or {
        head: Box<ExprPos>,
        rest: Vec<ExprPos>,
    },

    // (! <cond>)
    Not {
        cond: Box<ExprPos>,
    },

    // (bshl <expr> <length>)
    BShl {
        expr: Box<ExprPos>,
        length: Box<ExprPos>,
    },

    // (bshr <expr> <length>)
    BShr {
        expr: Box<ExprPos>,
        length: Box<ExprPos>,
    },

    // (CMP EXPR EXPR)
    Compare {
        cmp: CompareOperator,
        lhs: Box<ExprPos>,
        rhs: Box<ExprPos>,
    },

    // (<func> <args>...)
    // <expr>.(<func> <args>...)
    Call(Call),

    IfExpr(IfExpr),

    // (* EXPR)
    Dref(Box<ExprPos>),

    // (& EXPR)
    TakeAddr(Box<ExprPos>),
}

impl Expr {
    fn comptime_eval(&mut self, env: &Env) -> Result<(), Error> {
        match self {
            Expr::Int(_)
            | Expr::UInt(_)
            | Expr::I32(_)
            | Expr::U32(_)
            | Expr::I64(_)
            | Expr::U64(_)
            | Expr::Boolean(_)
            | Expr::String(_)
            | Expr::EnumGet { .. }
            | Expr::Char(_) => Ok(()),
            Expr::Variable(name) => {
                *self = env.find_const(name, &Pos::zero())?;
                Ok(())
            }
            Expr::Sizeof { ty } => match ty {
                Type::Unknown => todo!(),
                Type::Void => todo!(),
                Type::Char => todo!(),
                Type::Int => todo!(),
                Type::UInt => todo!(),
                Type::Int32 | Type::Uint32 => {
                    *self = Expr::UInt(4);
                    Ok(())
                }
                Type::Int64 | Type::Uint64 => {
                    *self = Expr::UInt(8);
                    Ok(())
                }
                Type::Boolean => {
                    *self = Expr::UInt(1);
                    Ok(())
                }
                _ => todo!(),
            },
            Expr::Compute { op, head, rest } => {
                match op {
                    Arithmetic::Add => {
                        head.expr.comptime_eval(env)?;

                        for each in rest.iter_mut() {
                            each.expr.comptime_eval(env)?;
                        }

                        let Expr::UInt(mut n) = head.expr else {
                            todo!()
                        };

                        for each in rest.iter() {
                            let Expr::UInt(x) = each.expr else { todo!() };
                            n += x;
                        }

                        *self = Expr::UInt(n);
                    }
                    Arithmetic::Mul => {
                        head.expr.comptime_eval(env)?;

                        for each in rest.iter_mut() {
                            each.expr.comptime_eval(env)?;
                        }

                        let Expr::UInt(mut n) = head.expr else {
                            todo!()
                        };

                        for each in rest.iter() {
                            let Expr::UInt(x) = each.expr else { todo!() };
                            n *= x;
                        }

                        *self = Expr::UInt(n);
                    }
                    Arithmetic::Sub => todo!(),
                    Arithmetic::Div => todo!(),
                    Arithmetic::Mod => todo!(),
                };
                Ok(())
            }
            _ => todo!(),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct ExprPos {
    pub expr: Expr,
    pub pos: Pos,
    pub ty: Type,
}

impl ExprPos {
    pub fn new(expr: Expr, pos: Pos) -> Self {
        Self {
            expr,
            pos,
            ty: Type::Unknown,
        }
    }

    fn infer_type(&mut self, env: &Env) -> Result<Type, Error> {
        let result_ty = match &mut self.expr {
            Expr::Null => Type::Null,
            Expr::Int(_) => Type::Int,
            Expr::UInt(_) => Type::UInt,
            Expr::I32(_) => Type::Int32,
            Expr::U32(_) => Type::Uint32,
            Expr::I64(_) => Type::Int64,
            Expr::U64(_) => Type::Uint64,
            Expr::String(_) => Type::Char.to_const().to_pointer(),
            Expr::Variable(name) => env.find_var_type(name, &self.pos)?,
            Expr::Boolean(_) => Type::Boolean,
            Expr::Sizeof { .. } => Type::UInt,
            Expr::Char(_) => Type::Char,
            Expr::Take { expr, index } => {
                let expr_ty = expr.infer_type(env)?;

                let index_ty = decay_const(index.infer_type(env)?);

                match index_ty {
                    Type::UInt => (),
                    _ => {
                        return Err(Error::Type {
                            pos: self.pos.clone(),
                            found: index_ty,
                            expect: Type::UInt,
                        });
                    }
                };

                match expr_ty {
                    Type::Pointer(inner) => *inner,
                    Type::Const(ref inner) => *match &**inner {
                        Type::Pointer(inner) => inner.clone(),
                        _ => {
                            return Err(Error::Syntax {
                                error: format!("can't use '[]' at type '{:?}'", expr_ty),
                                pos: self.pos.clone(),
                            })
                        }
                    },
                    Type::Array {
                        element_type,
                        length: _,
                    } => *element_type,
                    _ => {
                        return Err(Error::Syntax {
                            error: format!("can't use '[]' at type '{:?}'", expr_ty),
                            pos: self.pos.clone(),
                        })
                    }
                }
            }
            Expr::EnumGet { name, variant } => {
                let Ok(cenum) = env.find_enum(&name, &self.pos) else {
                    return Err(Error::Syntax {
                        error: format!("undefined enum {name}"),
                        pos: self.pos.clone(),
                    });
                };

                let mut result = None;

                for v in cenum.variants {
                    if v.inner == *variant {
                        result = Some(Type::Ref { name: name.clone() });
                    }
                }

                match result {
                    Some(some) => some,
                    None => {
                        return Err(Error::Syntax {
                            error: format!("no variant named {variant} in enum {name}"),
                            pos: self.pos.clone(),
                        })
                    }
                }
            }
            Expr::Get { expr, name } => {
                let ty = expr.infer_type(env)?;

                let (ty, is_const) = match ty {
                    Type::Const(inner) => (*inner, true),
                    _ => (ty, false),
                };

                let ty = match ty {
                    Type::Ref { name: ref_name, .. } => {
                        if let Ok(cstruct) = env.find_struct(&ref_name, &self.pos) {
                            let Some(field) = cstruct
                                .fields
                                .iter()
                                .find(|field| field.inner.name == *name)
                            else {
                                return Err(Error::Syntax {
                                    error: format!("no field name '{name}' of '{}'", cstruct.name),
                                    pos: expr.pos.clone(),
                                });
                            };
                            field.inner.ty.clone()
                        } else if let Ok(cunion) = env.find_union(&ref_name, &self.pos) {
                            let Some(variant) = cunion
                                .variants
                                .iter()
                                .find(|field| field.inner.name == *name)
                            else {
                                return Err(Error::Syntax {
                                    error: format!("no variant name '{name}' of '{}'", cunion.name),
                                    pos: expr.pos.clone(),
                                });
                            };
                            variant.inner.ty.clone()
                        } else {
                            return Err(Error::Syntax {
                                error: format!("{ref_name} is not a struct or a union"),
                                pos: expr.pos.clone(),
                            });
                        }
                    }
                    _ => {
                        return Err(Error::Syntax {
                            error: format!("invalid '.' of '{ty:?}'"),
                            pos: expr.pos.clone(),
                        })
                    }
                };

                if is_const {
                    Type::Const(Box::new(ty))
                } else {
                    ty
                }
            }
            Expr::Compare { cmp: _, lhs, rhs } => {
                check_binary_type(env, lhs, rhs)?;
                Type::Boolean
            }
            Expr::Not { cond } => {
                let ty = cond.infer_type(env)?;

                if !matches!(ty, Type::Boolean) {
                    return Err(Error::Type {
                        pos: cond.pos.clone(),
                        found: ty,
                        expect: Type::Boolean,
                    });
                }

                Type::Boolean
            }
            Expr::And { head, rest } | Expr::Or { head, rest } => {
                let ty = head.infer_type(env)?;

                if !matches!(ty, Type::Boolean) {
                    return Err(Error::Type {
                        pos: head.pos.clone(),
                        found: ty,
                        expect: Type::Boolean,
                    });
                }

                for x in rest {
                    let x_ty = x.infer_type(env)?;

                    if !matches!(x_ty, Type::Boolean) {
                        return Err(Error::Type {
                            pos: x.pos.clone(),
                            found: x_ty,
                            expect: Type::Boolean,
                        });
                    }
                }

                Type::Boolean
            }
            Expr::BShl { expr, length } | Expr::BShr { expr, length } => {
                // TODO: check expr_ty
                let expr_ty = expr.infer_type(env)?;

                let length_ty = length.infer_type(env)?;
                let Type::UInt = length_ty else {
                    return Err(Error::Type {
                        pos: length.pos.clone(),
                        found: length_ty,
                        expect: Type::UInt,
                    });
                };

                expr_ty
            }
            Expr::Compute { op, head, rest } => {
                let ty = head.infer_type(env)?;

                if matches!(op, Arithmetic::Add) && matches!(ty, Type::Pointer(_)) {
                    for x in rest {
                        let x_ty = x.infer_type(env)?;
                        if !matches!(x_ty, Type::UInt) {
                            return Err(Error::Type {
                                pos: x.pos.clone(),
                                found: x_ty,
                                expect: Type::UInt,
                            });
                        }
                    }
                } else {
                    for x in rest {
                        let x_ty = x.infer_type(env)?;
                        if !ty.can_assign(&x_ty) {
                            return Err(Error::Type {
                                pos: x.pos.clone(),
                                found: x_ty,
                                expect: ty,
                            });
                        }
                    }
                }

                ty
            }
            Expr::IfExpr(IfExpr { branch, otherwise }) => {
                let mut result_ty = None;

                for (cond, result) in branch {
                    let ty = cond.infer_type(env)?;

                    if !ty.can_assign(&Type::Boolean) {
                        return Err(Error::Type {
                            pos: cond.pos.clone(),
                            found: ty,
                            expect: Type::Boolean,
                        });
                    }

                    let ty = result.infer_type(env)?;

                    match &result_ty {
                        None => {
                            result_ty = Some(ty);
                        }
                        Some(result_ty) => {
                            if *result_ty != ty {
                                return Err(Error::Type {
                                    pos: result.pos.clone(),
                                    found: ty,
                                    expect: result_ty.clone(),
                                });
                            }
                        }
                    }
                }

                let else_ty = otherwise.infer_type(env)?;

                if let Some(result_ty) = result_ty {
                    if else_ty != result_ty {
                        return Err(Error::Type {
                            pos: otherwise.pos.clone(),
                            found: else_ty,
                            expect: result_ty,
                        });
                    }
                }

                else_ty
            }
            Expr::Call(call) => infer_func_type(&self.pos, call, env)?,
            Expr::Cast { expr, ty } => {
                expr.infer_type(env)?;

                ty.clone()
            }
            Expr::Dref(expr) => {
                let ty = expr.infer_type(env)?;

                let inner = match ty {
                    Type::Pointer(inner) => inner,
                    // TODO: make sure this is checked.
                    Type::OptionalPointer(inner) => inner,
                    _ => {
                        return Err(Error::CanNotDref {
                            ty,
                            pos: self.pos.clone(),
                        });
                    }
                };

                *inner
            }
            Expr::TakeAddr(expr) => {
                let ty = expr.infer_type(env)?;

                Type::Pointer(Box::new(ty))
            }
        };

        self.ty = result_ty.clone();

        Ok(result_ty)
    }

    fn infer_type_const(&mut self, env: &Env) -> Result<(), Error> {
        match &mut self.expr {
            Expr::Null => (),
            Expr::Cast { ty, expr } => {
                ty.comptime_eval(env)?;
                expr.infer_type_const(env)?;
            }
            Expr::Sizeof { ty } => {
                ty.comptime_eval(env)?;
            }
            Expr::Call(Call {
                is_macro: _,
                is_method: _,
                func: _,
                args,
            }) => {
                for arg in args {
                    arg.infer_type_const(env)?;
                }
            }
            Expr::IfExpr(IfExpr {
                branch,
                otherwise: default,
            }) => {
                for each in branch {
                    each.0.infer_type_const(env)?;
                    each.1.infer_type_const(env)?;
                }
                default.infer_type(env)?;
            }
            Expr::Dref(inner) => {
                inner.infer_type_const(env)?;
            }
            Expr::TakeAddr(inner) => {
                inner.infer_type_const(env)?;
            }
            Expr::Int(_) => (),
            Expr::UInt(_) => (),
            Expr::I32(_) => (),
            Expr::U32(_) => (),
            Expr::I64(_) => (),
            Expr::U64(_) => (),
            Expr::Boolean(_) => (),
            Expr::String(_) => (),
            Expr::Char(_) => (),
            Expr::Variable(_) => (),
            Expr::EnumGet { .. } => (),
            Expr::Get { expr, name: _ } => {
                expr.infer_type_const(env)?;
            }
            Expr::Take { expr, index } => {
                expr.infer_type_const(env)?;
                index.infer_type_const(env)?;
            }
            Expr::Compute { op: _, head, rest }
            | Expr::And { head, rest }
            | Expr::Or { head, rest } => {
                head.infer_type_const(env)?;
                for each in rest {
                    each.infer_type_const(env)?;
                }
            }
            Expr::BShl { expr, length } | Expr::BShr { expr, length } => {
                expr.infer_type_const(env)?;
                length.infer_type_const(env)?;
            }
            Expr::Not { cond } => {
                cond.infer_type_const(env)?;
            }
            Expr::Compare { cmp: _, lhs, rhs } => {
                lhs.infer_type_const(env)?;
                rhs.infer_type_const(env)?;
            }
        }
        Ok(())
    }
}

fn infer_func_type(pos: &Pos, call: &mut Call, env: &Env) -> Result<Type, Error> {
    let mut arg_types = Vec::with_capacity(call.args.len());

    for arg in &mut call.args {
        arg_types.push(arg.infer_type(env)?);
    }

    if call.func == "printf" {
        return Ok(Type::Int);
    }

    if call.is_macro {
        match call.func.as_ref() {
            "file" => return Ok(Type::Char.to_const().to_pointer()),
            "line" => return Ok(Type::UInt),
            "print" | "println" | "eprint" | "eprintln" => return Ok(Type::Void),
            _ => (),
        }
    }

    if call.is_method {
        let type_name = type_to_func_scope(&arg_types[0], pos)?;
        call.func = format!("{type_name}_{func}", func = call.func).into();
    }

    let ty = env.find_var_type(&call.func, &pos)?;

    let Type::Func {
        return_type,
        mut paras,
    } = ty
    else {
        return Err(Error::NotCallable {
            name: call.func.clone(),
            pos: pos.clone(),
            ty,
        });
    };

    if arg_types.len() != paras.len() {
        return Err(Error::ArgumentCount {
            pos: pos.clone(),
            found: arg_types.len(),
            expect: paras.len(),
        });
    }

    for (found, expect) in arg_types.iter_mut().zip(paras.iter_mut()) {
        if !expect.can_assign(found) {
            return Err(Error::Type {
                pos: pos.clone(),
                found: found.clone(),
                expect: expect.clone(),
            });
        }
    }

    Ok(*return_type)
}

fn type_to_func_scope(ty: &Type, pos: &Pos) -> Result<CowStr, Error> {
    Ok(match &ty {
        Type::Pointer(inner) => type_to_func_scope(inner.as_ref(), pos)?,
        Type::Const(inner) => type_to_func_scope(inner.as_ref(), pos)?,
        Type::Ref { name, .. } => name.clone(),
        _ => {
            return Err(Error::Syntax {
                error: "".to_string(),
                pos: pos.clone(),
            })
        }
    })
}

fn check_binary_type(
    env: &Env,
    lhs: &mut Box<ExprPos>,
    rhs: &mut Box<ExprPos>,
) -> Result<Type, Error> {
    let lhs_ty = lhs.infer_type(env)?;
    let rhs_ty = rhs.infer_type(env)?;

    if !(lhs_ty.can_assign(&rhs_ty) || rhs_ty.can_assign(&lhs_ty)) {
        return Err(Error::Type {
            pos: lhs.pos.clone(),
            found: lhs_ty.clone(),
            expect: rhs_ty.clone(),
        });
    }

    Ok(lhs_ty)
}

#[derive(Debug)]
pub struct Switch {
    pub head: Box<ExprPos>,
    pub branch: Vec<SwitchBranch>,
    pub has_default: bool,
    pub default: Vec<StatPos>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct IfExpr {
    pub branch: Vec<(ExprPos, ExprPos)>,
    pub otherwise: Box<ExprPos>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct Call {
    pub is_macro: bool,
    pub is_method: bool,
    pub func: CowStr,
    pub args: Vec<ExprPos>,
}

#[derive(Debug)]
pub struct SwitchBranch {
    pub variant: CowStr,
    pub stat: Vec<StatPos>,
}

#[derive(Debug)]
pub enum Stat {
    Import(CowStr),
    Return(Option<ExprPos>),
    Struct(Struct),
    Enum(Enum),
    Union(Union),
    Func(Func),
    Let(Let),
    LetUninit(VarDecl),
    Expr(ExprPos),
    // <target> = <value>
    Assign {
        target: ExprPos,
        value: ExprPos,
    },
    // <target> <op>= <value>
    ComputeAssign {
        op: Arithmetic,
        target: ExprPos,
        value: ExprPos,
    },
    // const <decl.name> = <expr>
    // const <decl.name>: <decl.ty> = <expr>
    ComptimeConstant {
        decl: VarDecl,
        expr: ExprPos,
        expr_display: String,
    },
    // if <head.cond> { <head.then> }
    // <elif <elif.cond> { <elif.then>... }>*
    // else { <otherwize>... }
    If {
        head: IfBranch,
        elif: Vec<IfBranch>,
        otherwise: Vec<StatPos>,
    },
    // while <cond> { <body>... }
    While {
        cond: ExprPos,
        body: Vec<StatPos>,
    },
    // for <index> = <iterator> { <body>... }
    For {
        index: CowStr,
        index_pos: Pos,
        index_ty: Type,

        iterator: ExprPos,
        body: Vec<StatPos>,
    },
    Break,
    Continue,
    Switch(Switch),
}

#[derive(Debug)]
pub struct IfBranch {
    pub cond: ExprPos,
    pub then: Vec<StatPos>,
}

impl IfBranch {
    fn infer_type_const(&mut self, env: &Env) -> Result<(), Error> {
        self.cond.infer_type_const(env)?;
        for each in &mut self.then {
            each.stat_infer_type_const(env)?;
        }
        Ok(())
    }
}

#[derive(Debug, PartialEq)]
pub struct Let {
    pub decl: VarDecl,
    pub expr: ExprPos,
}

#[derive(Debug, Clone, PartialEq)]
pub struct Struct {
    pub name: CowStr,
    pub fields: Vec<WithComment<VarDecl>>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct Enum {
    pub name: CowStr,
    pub variants: Vec<EnumVariant>,
}

pub type EnumVariant = WithComment<CowStr>;

#[derive(Debug, Clone, PartialEq)]
pub struct WithComment<T> {
    pub comment: Vec<String>,
    pub inner: T,
}

#[derive(Debug, Clone, PartialEq)]
pub struct Union {
    pub name: CowStr,
    pub variants: Vec<WithComment<VarDecl>>,
}

#[derive(Debug)]
pub struct Func {
    pub public: bool,
    pub return_type: Type,
    pub name: CowStr,
    pub paras: Vec<VarDecl>,
    pub body: Vec<StatPos>,
}

impl Func {
    fn parameter_types(&self) -> Vec<Type> {
        let mut tys = Vec::with_capacity(self.paras.len());

        for arg in self.paras.iter() {
            tys.push(arg.ty.clone())
        }

        tys
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct VarDecl {
    pub name: CowStr,
    pub ty: Type,
}

pub struct VarDeclC<'decl> {
    decl: &'decl VarDecl,
}

impl<'decl> VarDeclC<'decl> {
    pub fn new(decl: &'decl VarDecl) -> Self {
        Self { decl }
    }
}

impl<'decl> Display for VarDeclC<'decl> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.decl.ty {
            Type::Array {
                element_type,
                length,
            } => write!(
                f,
                "{} {}[{}]",
                element_type,
                self.decl.name,
                BareNumberC::new(&length.expr),
            ),
            Type::Pointer(inner) => write!(f, "{inner} *{name}", name = self.decl.name),
            _ => {
                write!(f, "{} {}", self.decl.ty, self.decl.name)
            }
        }
    }
}

#[derive(Debug, PartialEq, Clone)]
pub enum Type {
    Unknown,
    Void,
    Char,

    Int,
    UInt,
    Int32,
    Int64,
    Uint32,
    Uint64,
    Boolean,

    Pointer(Box<Type>),

    Null,
    OptionalPointer(Box<Type>),

    Const(Box<Type>),

    Array {
        element_type: Box<Type>,
        length: Box<ExprPos>,
    },

    Func {
        return_type: Box<Type>,
        paras: Vec<Type>,
    },

    Ref {
        name: CowStr,
    },
}

impl Type {
    fn comptime_eval(&mut self, env: &Env) -> Result<(), Error> {
        match self {
            Type::Array {
                element_type,
                length,
            } => {
                element_type.comptime_eval(env)?;
                length.infer_type(env)?;
            }
            Type::Func { return_type, paras } => {
                return_type.comptime_eval(env)?;
                for each in paras {
                    each.comptime_eval(env)?;
                }
            }
            Type::Pointer(inner) => {
                inner.comptime_eval(env)?;
            }
            Type::Const(inner) => {
                inner.comptime_eval(env)?;
            }
            _ => (),
        }
        Ok(())
    }

    pub fn to_pointer(self) -> Type {
        Type::Pointer(Box::new(self))
    }

    pub fn to_optional_pointer(self) -> Type {
        Type::OptionalPointer(Box::new(self))
    }

    pub fn to_const(self) -> Type {
        Type::Const(Box::new(self))
    }

    fn can_assign(&self, that: &Type) -> bool {
        let lhs = decay_const(self.clone());
        let rhs = decay_const(that.clone());

        if lhs == rhs {
            return true;
        }

        match (lhs, rhs) {
            (Type::OptionalPointer(_), Type::Null) => true,

            (Type::OptionalPointer(lhs) | Type::Pointer(lhs), Type::Pointer(rhs))
            | (Type::OptionalPointer(lhs), Type::OptionalPointer(rhs)) => {
                // Pointer to void is parent of all other pointers.
                if let Type::Void = *lhs {
                    true
                } else {
                    decay_const(*lhs) == *rhs
                }
            }
            _ => false,
        }
    }
}

impl Display for Type {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Type::Unknown => write!(f, "<UNKNOWN TYPE>"),
            Type::Void => write!(f, "void"),

            Type::Int => write!(f, "isize"),
            Type::UInt => write!(f, "usize"),
            Type::Int32 => write!(f, "i32"),
            Type::Int64 => write!(f, "i64"),
            Type::Uint32 => write!(f, "u32"),
            Type::Uint64 => write!(f, "u64"),

            Type::Boolean => write!(f, "bool"),
            Type::Pointer(ty) => write!(f, "{ty} *"),
            Type::OptionalPointer(ty) => write!(f, "{ty} *"),
            Type::Const(ty) => write!(f, "{ty} const"),
            Type::Char => write!(f, "char"),
            Type::Func { .. } => write!(f, "<{self:?}>"),
            Type::Ref { name, .. } => write!(f, "{}", name),
            Type::Null => write!(f, "null"),
            Type::Array {
                element_type,
                length,
            } => write!(
                f,
                "<array of {element_type}, length: {length}>",
                length = ExprC::new(&length.expr, 0)
            ),
        }
    }
}

#[derive(Debug)]
pub struct Code {
    stats: Vec<StatPos>,
}

impl Code {
    pub fn new(stats: Vec<StatPos>) -> Self {
        Self { stats }
    }

    pub fn bind_globals(&mut self, env: &mut Env) -> Result<(), Error> {
        let mut test_main_body = vec![];

        for stat in &self.stats {
            let name = match &stat.stat {
                Stat::Func(Func {
                    public: false,
                    return_type: Type::Void,
                    name,
                    paras,
                    body: _,
                }) => {
                    if name.starts_with("test_") && paras.is_empty() {
                        name
                    } else {
                        continue;
                    }
                }
                _ => continue,
            };

            test_main_body.push(StatPos {
                comments: vec![],
                stat: Stat::Expr(ExprPos {
                    expr: Expr::Call(Call {
                        is_macro: false,
                        is_method: false,
                        func: "puts".into(),
                        args: vec![ExprPos {
                            expr: Expr::String(format!("run test: {name}").into()),
                            pos: Pos::zero(),
                            ty: Type::Char.to_const().to_pointer(),
                        }],
                    }),
                    pos: Pos::zero(),
                    ty: Type::Void,
                }),
                pos: Pos::zero(),
            });

            test_main_body.push(StatPos {
                comments: vec![],
                stat: Stat::Expr(ExprPos {
                    expr: Expr::Call(Call {
                        is_macro: false,
                        is_method: false,
                        func: name.clone(),
                        args: vec![],
                    }),
                    pos: Pos::zero(),
                    ty: Type::Void,
                }),
                pos: Pos::zero(),
            })
        }

        if !test_main_body.is_empty() {
            let test_main = Func {
                public: false,
                return_type: Type::Void,
                name: "test_main".into(),
                paras: vec![],
                body: test_main_body,
            };

            self.stats.push(StatPos {
                comments: vec![],
                stat: Stat::Func(test_main),
                pos: Pos::zero(),
            });
        }

        for stat in &mut self.stats {
            stat.bind_globals(&env)?;
        }

        for stat in &mut self.stats {
            stat.refill_globals(&env)?;
        }

        Ok(())
    }

    pub fn infer_type(&mut self, env: &Env) -> Result<(), Error> {
        for stat in &mut self.stats {
            stat.stat_infer_type_const(&env)?;
        }

        for stat in &mut self.stats {
            stat.stat_infer_type_runtime(&env)?;
        }
        Ok(())
    }

    pub fn show(&self, gen_line: bool) {
        println!("/* Generated by makec compiler, DO NOT EDIT. */\n");

        let mut imports = HashSet::new();

        for stat in &self.stats {
            match &stat.stat {
                Stat::Import(path) => {
                    imports.insert(path.as_ref());

                    println!("#include <{path}>");
                }

                _ => (),
            }
        }

        let import_stdint = imports.contains("stdint.h");

        println!();

        struct IntegerName {
            c_name: &'static str,
            name: &'static str,
        }

        impl IntegerName {
            fn new(c_name: &'static str, name: &'static str) -> Self {
                Self { c_name, name }
            }

            fn typedef(&self) {
                println!("typedef {:<16} {};", self.c_name, self.name);
            }

            fn literal_macro(&self) {
                println!(
                    "#define {:<16} (({})X)",
                    format!("{}_(X)", self.name),
                    self.name
                );
            }
        }

        let stdints = [
            IntegerName::new("intptr_t", "isize"),
            IntegerName::new("uintptr_t", "usize"),
            IntegerName::new("int32_t", "i32"),
            IntegerName::new("uint32_t", "u32"),
            IntegerName::new("int64_t", "i64"),
            IntegerName::new("uint64_t", "u64"),
        ];

        if import_stdint {
            for each in &stdints {
                each.typedef();
            }
            for each in &stdints {
                each.literal_macro();
            }
        }

        for stat in &self.stats {
            print!("{}", CompimeConstant::new(&stat.stat));
        }

        for stat in &self.stats {
            print!("{}", CustomTypeDeclaration::new(&stat.stat));
        }

        for stat in &self.stats {
            print!("{}", CustomTypeDefinition::new(&stat));
        }

        for stat in &self.stats {
            print!("{}", FuncDeclare::new(&stat, 1));
        }

        for stat in &self.stats {
            print!("{}", StatPosC::new(stat, gen_line));
        }
    }
}

#[derive(Debug)]
pub struct StatPos {
    pub comments: Vec<String>,
    pub stat: Stat,
    pub pos: Pos,
}

impl StatPos {
    fn refill_globals(&mut self, env: &Env) -> Result<(), Error> {
        match &mut self.stat {
            Stat::Struct(cstruct) => {
                *cstruct = env.find_struct(&cstruct.name, &self.pos)?;
            }
            Stat::Enum(cenum) => {
                *cenum = env.find_enum(&cenum.name, &self.pos)?;
            }
            _ => (),
        }

        Ok(())
    }

    fn bind_globals(&mut self, env: &Env) -> Result<(), Error> {
        match &mut self.stat {
            Stat::Func(func) => {
                env.bind_var_type(
                    func.name.clone(),
                    Type::Func {
                        return_type: Box::new(func.return_type.clone()),
                        paras: func.parameter_types(),
                    },
                    &self.pos,
                )?;
            }
            Stat::Struct(cstruct) => {
                env.bind_struct(cstruct.name.clone(), cstruct.clone());
            }
            Stat::Enum(cenum) => {
                env.bind_enum(cenum.name.clone(), cenum.clone());
            }
            Stat::Union(c) => {
                env.bind_union(c.name.clone(), c.clone());
            }
            Stat::Import(name) => match name.as_ref() {
                "stdio.h" => load_stdio(&env)?,
                "assert.h" => load_assert(&env)?,
                "stdbool.h" => (),
                "stdint.h" => (),
                "inttypes.h" => (),
                "stdlib.h" => load_stdlib(&env)?,
                _ => panic!("undefined module: {name}"),
            },
            _ => (),
        }
        Ok(())
    }

    fn stat_infer_type_const(&mut self, env: &Env) -> Result<(), Error> {
        match &mut self.stat {
            Stat::ComptimeConstant { decl, expr, .. } => {
                expr.ty.comptime_eval(env)?;

                decl.ty = expr.infer_type(env)?;
                env.bind_var_type(decl.name.clone(), decl.ty.clone(), &self.pos)?;

                expr.expr.comptime_eval(env)?;
                env.bind_const(decl.name.clone(), expr.expr.clone(), &self.pos)?;
            }
            Stat::Return(expr) => {
                if let Some(expr) = expr {
                    expr.ty.comptime_eval(env)?;
                }
            }
            Stat::Struct(Struct { name: _, fields }) => {
                for field in fields {
                    field.inner.ty.comptime_eval(env)?;
                }
            }
            Stat::Union(Union { name: _, variants }) => {
                for variant in variants {
                    variant.inner.ty.comptime_eval(env)?;
                }
            }
            Stat::Func(Func {
                public: _,
                return_type,
                name: _,
                paras,
                body,
            }) => {
                return_type.comptime_eval(env)?;
                for para in paras {
                    para.ty.comptime_eval(env)?;
                }
                for each in body {
                    each.stat_infer_type_const(env)?;
                }
            }
            Stat::Let(Let { decl, expr: _ }) => {
                decl.ty.comptime_eval(env)?;
            }
            Stat::LetUninit(VarDecl { name: _, ty }) => {
                ty.comptime_eval(env)?;
            }
            Stat::Expr(expr) => {
                expr.infer_type_const(env)?;
            }
            Stat::Import(_) => (),
            Stat::Enum(_) => (),
            Stat::Assign { target, value } => {
                target.infer_type_const(env)?;
                value.infer_type_const(env)?;
            }
            Stat::ComputeAssign {
                op: _,
                target,
                value,
            } => {
                target.infer_type_const(env)?;
                value.infer_type_const(env)?;
            }
            Stat::If {
                head,
                elif,
                otherwise,
            } => {
                head.infer_type_const(env)?;
                for each in elif {
                    each.infer_type_const(env)?;
                }
                for each in otherwise {
                    each.stat_infer_type_const(env)?;
                }
            }
            Stat::While { cond, body } => {
                cond.infer_type_const(env)?;
                for each in body {
                    each.stat_infer_type_const(env)?;
                }
            }
            Stat::For {
                index: _,
                index_pos: _,
                index_ty,
                iterator,
                body,
            } => {
                index_ty.comptime_eval(env)?;
                iterator.infer_type_const(env)?;
                for each in body {
                    each.stat_infer_type_const(env)?;
                }
            }
            Stat::Break => (),
            Stat::Continue => (),
            Stat::Switch(Switch {
                head,
                branch,
                default,
                has_default: _,
            }) => {
                head.infer_type_const(env)?;

                for each in branch {
                    let env = Env::with_parent(env.clone());
                    head.infer_type_const(&env)?;
                    for each in &mut each.stat {
                        each.stat_infer_type_const(&env)?;
                    }
                }

                let env = Env::with_parent(env.clone());
                for each in default {
                    head.infer_type_const(&env)?;
                    each.stat_infer_type_const(&env)?;
                }
            }
        }
        Ok(())
    }

    fn stat_infer_type_runtime(&mut self, env: &Env) -> Result<(), Error> {
        match &mut self.stat {
            Stat::Import(_) => (),
            Stat::Return(expr) => {
                if let Some(expr) = expr {
                    expr.infer_type(env)?;
                }
            }
            Stat::Struct(c) => {
                *c = env.find_struct(&c.name, &self.pos)?;
            }
            Stat::Enum(c) => {
                *c = env.find_enum(&c.name, &self.pos)?;
            }
            Stat::Union(c) => {
                *c = env.find_union(&c.name, &self.pos)?;
            }
            Stat::Func(func) => {
                let env = Env::with_parent(env.clone());

                for arg in &mut func.paras {
                    env.bind_var_type(arg.name.clone(), arg.ty.clone(), &self.pos)?;
                }

                for stat in &mut func.body {
                    stat.stat_infer_type_runtime(&env)?;
                }
            }
            Stat::LetUninit(let_stat) => {
                env.bind_var_type(let_stat.name.clone(), let_stat.ty.clone(), &self.pos)?;
            }
            Stat::Let(let_stat) => {
                let expr_ty = let_stat.expr.infer_type(env)?;

                if let_stat.decl.ty == Type::Unknown {
                    let_stat.decl.ty = let_stat.expr.ty.clone();
                } else if !let_stat.decl.ty.can_assign(&expr_ty) {
                    return Err(Error::Type {
                        pos: self.pos.clone(),
                        found: expr_ty,
                        expect: let_stat.decl.ty.clone(),
                    });
                }

                let expr_ty = decay_const(expr_ty);

                env.bind_var_type(let_stat.decl.name.clone(), expr_ty, &self.pos)?;
            }
            Stat::Expr(expr) => {
                expr.infer_type(env)?;
            }
            Stat::Switch(Switch {
                head,
                branch,
                default,
                has_default,
            }) => {
                let head_ty = decay_const(head.infer_type(env)?);

                let mut expect_variants = HashSet::new();

                match &head_ty {
                    Type::Ref { name } => {
                        let head_enum = env.find_enum(&name, &head.pos)?;

                        for variant in &head_enum.variants {
                            expect_variants.insert(variant.inner.clone());
                        }
                    }
                    _ => {
                        return Err(Error::Syntax {
                            error: format!("can't support type `{}' on switch", head_ty),
                            pos: head.pos.clone(),
                        })
                    }
                }

                for branch in branch.iter_mut() {
                    let env = Env::with_parent(env.clone());

                    if !expect_variants.contains(&branch.variant) {
                        return Err(Error::Syntax {
                            error: format!(
                                "switch branch `{}' is not variant of `{}'",
                                branch.variant, head_ty,
                            ),
                            pos: head.pos.clone(),
                        });
                    }

                    for stat in &mut branch.stat {
                        stat.stat_infer_type_runtime(&env)?;
                    }
                }

                if *has_default {
                    let env = Env::with_parent(env.clone());
                    for stat in default {
                        stat.stat_infer_type_runtime(&env)?;
                    }
                } else {
                    if expect_variants.len() != branch.len() {
                        return Err(Error::Syntax {
                            error: format!(
                                "switch on `{}' need `{}' variants",
                                head_ty,
                                expect_variants.len(),
                            ),
                            pos: head.pos.clone(),
                        });
                    }
                }
            }
            Stat::While { cond, body } => {
                let env = Env::with_parent(env.clone());

                cond.infer_type(&env)?;
                for stat in body {
                    stat.stat_infer_type_runtime(&env)?;
                }
            }
            Stat::For {
                index,
                index_pos,
                iterator,
                body,
                index_ty,
            } => {
                let env = Env::with_parent(env.clone());

                let iter_ty = iterator.infer_type(&env)?;
                *index_ty = decay_const(index_type_of_iteraotor(iter_ty.clone()));

                env.bind_var_type(index.clone(), index_ty.clone(), &index_pos)?;

                for stat in body {
                    stat.stat_infer_type_runtime(&env)?;
                }
            }
            Stat::Break => (),
            Stat::Continue => (),
            Stat::If {
                head,
                elif,
                otherwise,
            } => {
                head.cond.infer_type(env)?;
                for stat in &mut head.then {
                    stat.stat_infer_type_runtime(env)?;
                }

                for branch in elif {
                    branch.cond.infer_type(env)?;
                    for stat in &mut branch.then {
                        stat.stat_infer_type_runtime(env)?;
                    }
                }

                for stat in otherwise {
                    stat.stat_infer_type_runtime(env)?;
                }
            }
            Stat::Assign { target, value } => {
                let ttype = target.infer_type(env)?;
                let vtype = value.infer_type(env)?;

                if !ttype.can_assign(&vtype) {
                    return Err(Error::Type {
                        pos: self.pos.clone(),
                        found: vtype,
                        expect: ttype,
                    });
                }
            }
            Stat::ComputeAssign {
                op: _,
                target,
                value,
            } => {
                let ttype = target.infer_type(env)?;
                let vtype = value.infer_type(env)?;

                if matches!(ttype, Type::Pointer(_)) {
                    if !matches!(vtype, Type::UInt) {
                        return Err(Error::Type {
                            pos: self.pos.clone(),
                            found: vtype,
                            expect: Type::UInt,
                        });
                    }
                } else {
                    if !ttype.can_assign(&vtype) {
                        return Err(Error::Type {
                            pos: self.pos.clone(),
                            found: vtype,
                            expect: ttype,
                        });
                    }
                }
            }
            Stat::ComptimeConstant { .. } => (),
        }
        Ok(())
    }
}

fn decay_const(ty: Type) -> Type {
    match ty {
        Type::Const(inner) => *inner,
        _ => ty,
    }
}

fn index_type_of_iteraotor(ty: Type) -> Type {
    ty
}
